//      motif_test.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 

#include <boost/foreach.hpp> 
#include <map>
#include "stl-helper.h"

#include "../gene_processing.h"
#include "../gene_em.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0);

void read_test();
void entropy_test();
void gibbs_test();
void motif_stats_test();
void em_test();
void em_max_test();
void em2_test();
void em2_max_test();

void extract_upstream();
void new_motif_stats_test();
void produce_score();

int main(int argc, char** argv)
{
  MSG("Motif alignment ... ");

  process_options(argc, argv);

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if ( vm.count("em") || vm.count("em-max") ) {
    int seed = vm["seed"].as<int>();
    int time_val = time(NULL), seed_val = 1;
    for (int i = 0; i < seed; i++)
      seed_val *= time_val;
    srand(seed_val);
  }
    
  if (vm.count("help"))
    usage();
  else if (vm.count("read"))
    read_test();
  else if (vm.count("entropy"))
    entropy_test();
  else if (vm.count("gibbs"))
    gibbs_test();
  else if (vm.count("mtstats"))
    new_motif_stats_test();
  else if (vm.count("em"))
    em_test();
  else if (vm.count("em-max"))
    em_max_test();
  //  else if (vm.count("em2"))
  //    em2_test();
  //  else if (vm.count("em2-max"))
  //    em2_max_test();
  else if (vm.count("upstream"))
    extract_upstream();
  else if (vm.count("score"))
    produce_score();
  else
    usage();

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  MSG("Duration = " << duration);
}

void produce_score()
{
  std::string inputFile = vm["input"].as<std::string>();
  std::string outputFile = vm["output"].as<std::string>();
  std::string pwmFile = vm["pwm"].as<std::string>();

  Sequences seqs;
  GroupSequences gseq;
  
  seqs.read_sequence(inputFile);
  gseq.group_sequences(seqs);
  
  int count = 0;
  for (int i = 0; i < sz(gseq); i++) count += sz(gseq[i]);
  MSG("procedure score" << " n_seq = " << sz(seqs) << " n_group = " << sz(gseq) << " n_seq = " << count);  
  
  PositionWeight pwm;
  pwm.read_pwm(pwmFile);
  MSG(pwm.toString());
  
  std::ofstream out;
  FileAccess::open2write(out, outputFile);
  
  count = 0;
  BOOST_FOREACH(const Sequences& g, gseq) {
    vd_t g_score(sz(g));
    int j = 0;
    BOOST_FOREACH(const std::string& seq, g) {
      double score = pwm.score(seq);
      g_score[j++] = score;
    }

    double max = arma::max(g_score);
    for (int i = 0; i < sz(g); i++) {
      out << g.annotation_[i] << (g_score[i] == max ? " MAX SCORE " : "" ) << ENDL 
          << g_score[i] << ENDL;
    }
    
    if (g_score[0] == max) count++;
  }
  out << "Correct = " << count << " / " << sz(gseq) << " accuracy = " << (double)(count) / sz(gseq) << ENDL;
  out.close();
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("input", po::value<std::string>(), "input file")
    ("output", po::value<std::string>(), "output file")
    ("annotation", po::value<std::string>(), "annotation file")
    ("pwm", po::value<std::string>(), "position wright matrix (PWM) file")
    ("seed", po::value<int>(), "random seed (REQUIRED)")

    // PARAMETERS
    ("pattern", po::value<std::string>(), "pattern")
    ("mtlen", po::value<int>()->default_value(6), "motif length")
    ("iter", po::value<int>()->default_value(1), "maximum number of iterations")
    ("run", po::value<int>()->default_value(10), "number of runs")
    ("simple", "simple annotation file (ORF only)")
    ("seqlen", po::value<int>()->default_value(30), "sequence length to extract")

    //JOB TYPE
    ("read", "read genome and annotation")
    ("entropy", "generate pattern and entropy of counts")
    ("gibbs", "gibbs sampling for motif finding")
    ("mtstats", "all statistics of a motif \n   --input --output --mtlen")
    ("em", "EM algorithm for motif\n   --input --run --iter --mtlen --seed --output")
    //("em2", "EM algorithm (2nd markov background) for motif\n   --input --run --iter --mtlen --output")
    ("em-max", "EM-MAX algorithm for motif\n   --input --run --iter --mtlen --seed --output")
    //("em2-max", "EM-MAX algorithm (2nd markov background) for motif\n   --input --run --iter --mtlen --output")
    ("upstream", "extract upstream sequences \n   --input --annotation --seqlen --output")
    ("score", "calculate score using PWM file \n   --input --pwm --output")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void usage(int exitcode)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}
